{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "tf_lenet_mnist.ipynb",
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/intel-analytics/BigDL/blob/main/python/orca/colab-notebook/quickstart/tf_lenet_mnist.ipynb)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xgV_QLGE9Lox"
      },
      "source": [
        "\r\n",
        "![image.png]()\n",
        "---"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fDBPZ0_rfBmU"
      },
      "source": [
        "##### Copyright 2016 The BigDL Authors."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xBWVU_bhfkY7"
      },
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\r\n",
        "# you may not use this file except in compliance with the License.\r\n",
        "# You may obtain a copy of the License at\r\n",
        "#\r\n",
        "#     http://www.apache.org/licenses/LICENSE-2.0\r\n",
        "#\r\n",
        "# Unless required by applicable law or agreed to in writing, software\r\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\r\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\r\n",
        "# See the License for the specific language governing permissions and\r\n",
        "# limitations under the License.\r\n",
        "#"
      ],
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "voMBntim9bMf"
      },
      "source": [
        "## **Environment Preparation**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ccXnsv_pMCHQ"
      },
      "source": [
        "**Install Java 8**\r\n",
        "\r\n",
        "Run the cell on the **Google Colab** to install jdk 1.8.\r\n",
        "\r\n",
        "**Note:** if you run this notebook on your computer, root permission is required when running the cell to install Java 8. (You may ignore this cell if Java 8 has already been set up in your computer).\r\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HWx_0DSnXVhz"
      },
      "source": [
        "# Install jdk8\r\n",
        "!apt-get install openjdk-8-jdk-headless -qq > /dev/null\r\n",
        "import os\r\n",
        "# Set environment variable JAVA_HOME.\r\n",
        "os.environ[\"JAVA_HOME\"] = \"/usr/lib/jvm/java-8-openjdk-amd64\"\r\n",
        "!update-alternatives --set java /usr/lib/jvm/java-8-openjdk-amd64/jre/bin/java\r\n",
        "!java -version"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "I_OS4HKJMNpv"
      },
      "source": [
        "**Install BigDL Orca**\r\n",
        "\r\n",
        "You can install the latest pre-release version using `pip install --pre --upgrade bigdl-orca`. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3qfT8CaC51hI"
      },
      "source": [
        "# Install latest pre-release version of BigDL Orca \r\n",
        "# Installing BigDL Orca from pip will automatically install pyspark, bigdl, and their dependencies.\r\n",
        "!pip install --pre --upgrade bigdl-orca"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uWQaFJBf55LJ"
      },
      "source": [
        "# Install required dependencies\r\n",
        "!pip install tensorflow==1.15.0 tensorflow-datasets==2.1.0"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "l-l4vel5N3qP"
      },
      "source": [
        "## **Distributed TensorFlow (v1.15) using Orca APIs**\r\n",
        "\r\n",
        "In this guide we will describe how to scale out TensorFlow (v1.15) programs using Orca in 4 simple steps."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RDPt2XTsT-Gv"
      },
      "source": [
        "# import necesary libraries and modules\r\n",
        "import argparse\r\n",
        "\r\n",
        "from bigdl.orca.learn.tf.estimator import Estimator\r\n",
        "from bigdl.orca import init_orca_context, stop_orca_context\r\n",
        "from bigdl.orca import OrcaContext"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zBdeoZzLWWlY"
      },
      "source": [
        "### **Step 1: Init Orca Context** "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QAQp0FcUWaH3"
      },
      "source": [
        "OrcaContext.log_output = True # recommended to set it to True when running BigDL in Jupyter notebook (this will display terminal's stdout and stderr in the Jupyter notebook).\r\n",
        "\r\n",
        "cluster_mode = \"local\"\r\n",
        "\r\n",
        "if cluster_mode == \"local\":  \r\n",
        "    init_orca_context(cluster_mode=\"local\", cores=4) # run in local mode\r\n",
        "    dataset_dir = \"~/tensorflow_datasets\"\r\n",
        "elif cluster_mode == \"k8s\":  \r\n",
        "    init_orca_context(cluster_mode=\"k8s\", num_nodes=2, cores=2) # run on K8s cluster\r\n",
        "elif cluster_mode == \"yarn\":  \r\n",
        "    init_orca_context(cluster_mode=\"yarn-client\", num_nodes=2, cores=2) # run on Hadoop YARN cluster\r\n",
        "    dataset_dir = \"hdfs:///tensorflow_datasets\""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9VyfhD6HVQpm"
      },
      "source": [
        "This is the only place where you need to specify local or distributed mode. View [Orca Context](https://bigdl.readthedocs.io/en/latest/doc/Orca/Overview/orca-context.html) for more details.\r\n",
        "\n",
        "**Note**: You should export HADOOP_CONF_DIR=/path/to/hadoop/conf/dir when you run on Hadoop YARN cluster. To use tensorflow_datasets on HDFS, you should correctly set HADOOP_HOME, HADOOP_HDFS_HOME, LD_LIBRARY_PATH, etc. For more details, please refer to TensorFlow documentation [link](https://github.com/tensorflow/docs/blob/r1.11/site/en/deploy/hadoop.md).\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ld-pV8e-Vb7D"
      },
      "source": [
        "### **Step 2: Define the Model**\r\n",
        "\r\n",
        "You may define your model, loss and metrics in the same way as in any standard (single node) TensorFlow program."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DC1P7_gP9FKm"
      },
      "source": [
        "import tensorflow as tf\r\n",
        "\r\n",
        "def accuracy(logits, labels):\r\n",
        "  predictions = tf.argmax(logits, axis=1, output_type=labels.dtype)\r\n",
        "  is_correct = tf.cast(tf.equal(predictions, labels), dtype=tf.float32)\r\n",
        "  return tf.reduce_mean(is_correct)\r\n",
        "\r\n",
        "\r\n",
        "def lenet(images):\r\n",
        "  with tf.variable_scope('LeNet', [images]):\r\n",
        "    net = tf.layers.conv2d(images, 32, (5, 5), activation=tf.nn.relu, name='conv1')\r\n",
        "    net = tf.layers.max_pooling2d(net, (2, 2), 2, name='pool1')\r\n",
        "    net = tf.layers.conv2d(net, 64, (5, 5), activation=tf.nn.relu, name='conv2')\r\n",
        "    net = tf.layers.max_pooling2d(net, (2, 2), 2, name='pool2')\r\n",
        "    net = tf.layers.flatten(net)\r\n",
        "    net = tf.layers.dense(net, 1024, activation=tf.nn.relu, name='fc3')\r\n",
        "    logits = tf.layers.dense(net, 10)\r\n",
        "    return logits\r\n",
        "\r\n",
        "# tensorflow inputs\r\n",
        "images = tf.placeholder(dtype=tf.float32, shape=(None, 28, 28, 1))\r\n",
        "# tensorflow labels\r\n",
        "labels = tf.placeholder(dtype=tf.int32, shape=(None,))\r\n",
        "\r\n",
        "logits = lenet(images)\r\n",
        "\r\n",
        "loss = tf.reduce_mean(tf.losses.sparse_softmax_cross_entropy(logits=logits, labels=labels))\r\n",
        "\r\n",
        "acc = accuracy(logits, labels)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CiTnHATGYn_f"
      },
      "source": [
        "### **Step 3: Define Train Dataset**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-YHW93myYykE"
      },
      "source": [
        "import tensorflow_datasets as tfds\r\n",
        "\r\n",
        "def preprocess(data):\r\n",
        "  data['image'] = tf.cast(data[\"image\"], tf.float32) / 255.\r\n",
        "  return data['image'], data['label']"
      ],
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yMyUkr0HYzNV"
      },
      "source": [
        "# get DataSet\r\n",
        "mnist_train = tfds.load(name=\"mnist\", split=\"train\", data_dir=dataset_dir)\r\n",
        "mnist_test = tfds.load(name=\"mnist\", split=\"test\", data_dir=dataset_dir)\r\n",
        "\r\n",
        "mnist_train = mnist_train.map(preprocess)\r\n",
        "mnist_test = mnist_test.map(preprocess)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CnbOmWHsXOmg"
      },
      "source": [
        "### **Step 4: Fit with Orca Estimator**\r\n",
        "\r\n",
        "First, create an Estimator.\r\n",
        "\r\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bsNbQROFXNgU"
      },
      "source": [
        "est = Estimator.from_graph(inputs=images,\r\n",
        "                           outputs=logits,\r\n",
        "                           labels=labels,\r\n",
        "                           loss=loss,\r\n",
        "                           optimizer=tf.train.AdamOptimizer(),\r\n",
        "                           metrics={\"acc\": acc})"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7bNtls5WXnzU"
      },
      "source": [
        "Next, fit the Estimator."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "imhbbDNXXsXR"
      },
      "source": [
        "max_epoch = 1\r\n",
        "est.fit(data=mnist_train,\r\n",
        "        batch_size=320,\r\n",
        "        epochs=max_epoch,\r\n",
        "        validation_data=mnist_test)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vRdASNtRAbiN"
      },
      "source": [
        "Finally, evaluate using the Estimator."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GDNoiHiquX0A"
      },
      "source": [
        "result = est.evaluate(mnist_test)\r\n",
        "print(result)\r\n",
        "\r\n",
        "est.save_tf_checkpoint(\"/tmp/lenet/model\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yyRtP6lqieej"
      },
      "source": [
        "Now, the accuracy of this model has reached 98%."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nC4KhlHxXakk"
      },
      "source": [
        "# Stop orca context when your program finishes\r\n",
        "stop_orca_context()"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}
